# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.

# RUN: pyasc -e codegen --args x=*float32 y=*float32 z=*float32 --constexprs BLOCK_LENGTH=512 BUFFER_NUM=1 TILE_LENGTH=64 TILE_NUM=8 --kernel vadd_kernel %s | FileCheck %s
# CHECK:      #loc37 = loc("i"(#loc1))
# CHECK-NEXT: #loc38 = loc("z_gm"(#loc1))
# CHECK-NEXT: #loc39 = loc("out_queue_z"(#loc1))
# CHECK-NEXT: #loc40 = loc("z_gm"(#loc6))
# CHECK-NEXT: #loc41 = loc("in_queue_x"(#loc6))
# CHECK-NEXT: #loc42 = loc("in_queue_y"(#loc6))
# CHECK-NEXT: #loc43 = loc("out_queue_z"(#loc6))
# CHECK-NEXT: #loc44 = loc("i"(#loc14))
# CHECK-NEXT: #loc45 = loc("x_gm"(#loc14))
# CHECK-NEXT: #loc46 = loc("y_gm"(#loc14))
# CHECK-NEXT: #loc47 = loc("in_queue_x"(#loc14))
# CHECK-NEXT: #loc48 = loc("in_queue_y"(#loc14))
# CHECK-NEXT: #loc49 = loc("x"(#loc23))
# CHECK-NEXT: #loc50 = loc("y"(#loc23))
# CHECK-NEXT: #loc51 = loc("z"(#loc23))
# CHECK-NEXT: module {
# CHECK-NEXT:   func.func @copy_out(%arg0: i32 loc("i"(#loc1)), %arg1: !ascendc.global_tensor<*xf32> loc("z_gm"(#loc1)), %arg2: !ascendc.queue<vec_out, 1> loc("out_queue_z"(#loc1))) attributes {ascendc.aicore} {
# CHECK-NEXT:     %0 = ascendc.que_bind.deque_tensor %arg2 : !ascendc.queue<vec_out, 1>, !ascendc.local_tensor<*xf32> loc(#loc2)
# CHECK-NEXT:     %c64_i32 = arith.constant 64 : i32 loc(#loc3)
# CHECK-NEXT:     %1 = arith.muli %arg0, %c64_i32 : i32 loc(#loc3)
# CHECK-NEXT:     %2 = ascendc.global_tensor.subindex %arg1[%1] : !ascendc.global_tensor<*xf32>, i32, !ascendc.global_tensor<*xf32> loc(#loc3)
# CHECK-NEXT:     %c64_i32_0 = arith.constant 64 : i32 loc(#loc4)
# CHECK-NEXT:     ascendc.data_copy_l2 %2, %0, %c64_i32_0 : !ascendc.global_tensor<*xf32>, !ascendc.local_tensor<*xf32>, i32 loc(#loc4)
# CHECK-NEXT:     ascendc.que_bind.free_tensor %arg2, %0 : !ascendc.queue<vec_out, 1>, !ascendc.local_tensor<*xf32> loc(#loc5)
# CHECK-NEXT:     return loc(#loc5)
# CHECK-NEXT:   } loc(#loc1)
# CHECK-NEXT:   func.func @compute(%arg0: !ascendc.global_tensor<*xf32> loc("z_gm"(#loc6)), %arg1: !ascendc.queue<vec_in, 1> loc("in_queue_x"(#loc6)), %arg2: !ascendc.queue<vec_in, 1> loc("in_queue_y"(#loc6)), %arg3: !ascendc.queue<vec_out, 1> loc("out_queue_z"(#loc6))) attributes {ascendc.aicore} {
# CHECK-NEXT:     %0 = ascendc.que_bind.deque_tensor %arg1 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc7)
# CHECK-NEXT:     %1 = ascendc.que_bind.deque_tensor %arg2 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc8)
# CHECK-NEXT:     %2 = ascendc.que_bind.alloc_tensor %arg3 : !ascendc.queue<vec_out, 1>, !ascendc.local_tensor<*xf32> loc(#loc9)
# CHECK-NEXT:     %c64_i32 = arith.constant 64 : i32 loc(#loc10)
# CHECK-NEXT:     ascendc.add_l2 %2, %0, %1, %c64_i32 : !ascendc.local_tensor<*xf32>, !ascendc.local_tensor<*xf32>, !ascendc.local_tensor<*xf32>, i32 loc(#loc10)
# CHECK-NEXT:     ascendc.que_bind.enque_tensor %arg3, %2 : !ascendc.queue<vec_out, 1>, !ascendc.local_tensor<*xf32> loc(#loc11)
# CHECK-NEXT:     ascendc.que_bind.free_tensor %arg1, %0 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc12)
# CHECK-NEXT:     ascendc.que_bind.free_tensor %arg2, %1 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc13)
# CHECK-NEXT:     return loc(#loc13)
# CHECK-NEXT:   } loc(#loc6)
# CHECK-NEXT:   func.func @copy_in(%arg0: i32 loc("i"(#loc14)), %arg1: !ascendc.global_tensor<*xf32> loc("x_gm"(#loc14)), %arg2: !ascendc.global_tensor<*xf32> loc("y_gm"(#loc14)), %arg3: !ascendc.queue<vec_in, 1> loc("in_queue_x"(#loc14)), %arg4: !ascendc.queue<vec_in, 1> loc("in_queue_y"(#loc14))) attributes {ascendc.aicore} {
# CHECK-NEXT:     %0 = ascendc.que_bind.alloc_tensor %arg3 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc15)
# CHECK-NEXT:     %1 = ascendc.que_bind.alloc_tensor %arg4 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc16)
# CHECK-NEXT:     %c64_i32 = arith.constant 64 : i32 loc(#loc17)
# CHECK-NEXT:     %2 = arith.muli %arg0, %c64_i32 : i32 loc(#loc17)
# CHECK-NEXT:     %3 = ascendc.global_tensor.subindex %arg1[%2] : !ascendc.global_tensor<*xf32>, i32, !ascendc.global_tensor<*xf32> loc(#loc17)
# CHECK-NEXT:     %c64_i32_0 = arith.constant 64 : i32 loc(#loc18)
# CHECK-NEXT:     ascendc.data_copy_l2 %0, %3, %c64_i32_0 : !ascendc.local_tensor<*xf32>, !ascendc.global_tensor<*xf32>, i32 loc(#loc18)
# CHECK-NEXT:     %c64_i32_1 = arith.constant 64 : i32 loc(#loc19)
# CHECK-NEXT:     %4 = arith.muli %arg0, %c64_i32_1 : i32 loc(#loc19)
# CHECK-NEXT:     %5 = ascendc.global_tensor.subindex %arg2[%4] : !ascendc.global_tensor<*xf32>, i32, !ascendc.global_tensor<*xf32> loc(#loc19)
# CHECK-NEXT:     %c64_i32_2 = arith.constant 64 : i32 loc(#loc20)
# CHECK-NEXT:     ascendc.data_copy_l2 %1, %5, %c64_i32_2 : !ascendc.local_tensor<*xf32>, !ascendc.global_tensor<*xf32>, i32 loc(#loc20)
# CHECK-NEXT:     ascendc.que_bind.enque_tensor %arg3, %0 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc21)
# CHECK-NEXT:     ascendc.que_bind.enque_tensor %arg4, %1 : !ascendc.queue<vec_in, 1>, !ascendc.local_tensor<*xf32> loc(#loc22)
# CHECK-NEXT:     return loc(#loc22)
# CHECK-NEXT:   } loc(#loc14)
# CHECK-NEXT:   func.func @vadd_kernel(%arg0: memref<?xf32, 22> loc("x"(#loc23)), %arg1: memref<?xf32, 22> loc("y"(#loc23)), %arg2: memref<?xf32, 22> loc("z"(#loc23))) attributes {ascendc.aicore, ascendc.global} {
# CHECK-NEXT:     %0 = ascendc.get_block_idx : i32 loc(#loc24)
# CHECK-NEXT:     %c512_i32 = arith.constant 512 : i32 loc(#loc25)
# CHECK-NEXT:     %1 = arith.muli %0, %c512_i32 : i32 loc(#loc25)
# CHECK-NEXT:     %2 = arith.index_cast %1 : i32 to index loc(#loc26)
# CHECK-NEXT:     %3 = emitasc.ptr_offset %arg0[%2] : memref<?xf32, 22>, memref<?xf32, 22> loc(#loc26)
# CHECK-NEXT:     %4 = ascendc.global_tensor : !ascendc.global_tensor<*xf32> loc(#loc26)
# CHECK-NEXT:     ascendc.global_tensor.set_global_buffer %4, %3 : !ascendc.global_tensor<*xf32>, memref<?xf32, 22> loc(#loc26)
# CHECK-NEXT:     %5 = arith.index_cast %1 : i32 to index loc(#loc27)
# CHECK-NEXT:     %6 = emitasc.ptr_offset %arg1[%5] : memref<?xf32, 22>, memref<?xf32, 22> loc(#loc27)
# CHECK-NEXT:     %7 = ascendc.global_tensor : !ascendc.global_tensor<*xf32> loc(#loc27)
# CHECK-NEXT:     ascendc.global_tensor.set_global_buffer %7, %6 : !ascendc.global_tensor<*xf32>, memref<?xf32, 22> loc(#loc27)
# CHECK-NEXT:     %8 = arith.index_cast %1 : i32 to index loc(#loc28)
# CHECK-NEXT:     %9 = emitasc.ptr_offset %arg2[%8] : memref<?xf32, 22>, memref<?xf32, 22> loc(#loc28)
# CHECK-NEXT:     %10 = ascendc.global_tensor : !ascendc.global_tensor<*xf32> loc(#loc28)
# CHECK-NEXT:     ascendc.global_tensor.set_global_buffer %10, %9 : !ascendc.global_tensor<*xf32>, memref<?xf32, 22> loc(#loc28)
# CHECK-NEXT:     %11 = ascendc.pipe loc(#loc29)
# CHECK-NEXT:     %12 = ascendc.queue : <vec_in, 1> loc(#loc30)
# CHECK-NEXT:     %13 = ascendc.queue : <vec_in, 1> loc(#loc31)
# CHECK-NEXT:     %14 = ascendc.queue : <vec_out, 1> loc(#loc32)
# CHECK-NEXT:     %c1_i32 = arith.constant 1 : i32 loc(#loc33)
# CHECK-NEXT:     %c256_i32 = arith.constant 256 : i32 loc(#loc33)
# CHECK-NEXT:     ascendc.pipe.init_queue %11, %12, %c1_i32, %c256_i32 : !ascendc.queue<vec_in, 1>, i32, i32 loc(#loc33)
# CHECK-NEXT:     %c1_i32_0 = arith.constant 1 : i32 loc(#loc34)
# CHECK-NEXT:     %c256_i32_1 = arith.constant 256 : i32 loc(#loc34)
# CHECK-NEXT:     ascendc.pipe.init_queue %11, %13, %c1_i32_0, %c256_i32_1 : !ascendc.queue<vec_in, 1>, i32, i32 loc(#loc34)
# CHECK-NEXT:     %c1_i32_2 = arith.constant 1 : i32 loc(#loc35)
# CHECK-NEXT:     %c256_i32_3 = arith.constant 256 : i32 loc(#loc35)
# CHECK-NEXT:     ascendc.pipe.init_queue %11, %14, %c1_i32_2, %c256_i32_3 : !ascendc.queue<vec_out, 1>, i32, i32 loc(#loc35)
# CHECK-NEXT:     %c0_i32 = arith.constant 0 : i32 loc(#loc36)
# CHECK-NEXT:     %c8_i32 = arith.constant 8 : i32 loc(#loc36)
# CHECK-NEXT:     %c1_i32_4 = arith.constant 1 : i32 loc(#loc36)
# CHECK-NEXT:     scf.for %arg3 = %c0_i32 to %c8_i32 step %c1_i32_4  : i32 {
# CHECK-NEXT:       func.call @copy_in(%arg3, %4, %7, %12, %13) : (i32, !ascendc.global_tensor<*xf32>, !ascendc.global_tensor<*xf32>, !ascendc.queue<vec_in, 1>, !ascendc.queue<vec_in, 1>) -> () loc(#loc)
# CHECK-NEXT:       func.call @compute(%10, %12, %13, %14) : (!ascendc.global_tensor<*xf32>, !ascendc.queue<vec_in, 1>, !ascendc.queue<vec_in, 1>, !ascendc.queue<vec_out, 1>) -> () loc(#loc)
# CHECK-NEXT:       func.call @copy_out(%arg3, %10, %14) : (i32, !ascendc.global_tensor<*xf32>, !ascendc.queue<vec_out, 1>) -> () loc(#loc)
# CHECK-NEXT:     } loc(#loc36)
# CHECK-NEXT:     return loc(#loc36)
# CHECK-NEXT:   } loc(#loc23)
# CHECK-NEXT: } loc(#loc)
# CHECK-NEXT: #loc = loc(unknown)

import os
import sys

root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
sys.path.append(root_dir)
THIS_FILE_NAME = __file__
FILE_PATH = os.path.dirname(os.path.realpath(THIS_FILE_NAME))
sys.path.append(os.path.join(FILE_PATH, "../../../python/test/kernels"))

from python.test.kernels.test_vadd import vadd_kernel
